package com.tuniu.agents.impl.base;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.util.Assert;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.agent.AgentOptions;
import com.tuniu.agents.agent.OptionsAgent;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.CustomerBeanOutputConverter;
import com.tuniu.agents.util.JsonUtil;

import reactor.core.publisher.Flux;

public class Dispatcher extends OptionsAgent {

    public static final String AGENT_ID = "Dispatcher";

    public static final String PROMPT_ID = "dispatcher";

    private final AgentManager agentManager;

    private final Map<String, String> selectionAgents;

    /**
     * {{user_message_current}} is assembled in Selection before coming here
     * Analyst:
     * User's previous requirements:
     * 1: XXX
     * 2: XXX
     * Current user requirement is:
     * XXX
     * Current recommended result is:
     * XXX (MD format)
     */
    public static final String MESSAGE_TEMPLATE = """
            {{user_message_his}}
            Current user requirement is:
            {{user_message_current}}
            """;

    public Dispatcher(AgentManagerGroup agentManagerGroup, Map<String, String> selectionAgents) {
        this(agentManagerGroup, AGENT_ID, PROMPT_ID, selectionAgents);
    }

    public Dispatcher(AgentManagerGroup agentManagerGroup, String agentId, String promptTemplateId, Map<String, String> selectionAgents) {
        super(agentManagerGroup.getModelBuilder(),
                agentManagerGroup.getChatMemory(),
                agentManagerGroup.getAdvisors().stream().map(advisor -> advisor instanceof AgentChatMemoryAdvisor ? new AgentChatMemoryAdvisor(agentManagerGroup.getChatMemory(), false, false) : advisor).collect(Collectors.toList()),
                agentManagerGroup.getPromptManager(),
                agentManagerGroup.getChatEnvironmentMemory(),
                new AgentOptions(agentId, promptTemplateId));
        this.agentOptions.getAdvisorParam().put(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10); // Only retrieve the most recent 10 conversation turns
        this.agentManager = agentManagerGroup.getAgentManager();
        this.selectionAgents = selectionAgents;
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> context) {
        int index = 1;
        StringBuffer agentSb = new StringBuffer();
        for(Map.Entry<String, String> entry : selectionAgents.entrySet()) {
            agentSb.append(index++).append(".").append(entry.getKey()).append(":").append(entry.getValue()).append("\n");
        }
        Dispatcher.DispatcherOut dispatcherOut =
                chatClient.prompt()
                        .system(s -> s.params(chatEnvironmentMemory.get(conversationId)))
                        .messages(messages)
                        .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId).
                                param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, agentOptions.getAdvisorParam().get(CHAT_MEMORY_RETRIEVE_SIZE_KEY)).
                                param(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY, this.agentId)
                                .param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE, MESSAGE_TEMPLATE)
                                .param(AgentsConstants.AgentsAdvisorKey.ASSISTANT_MESSAGE_HISTORY_NEED, "false")
                                .param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_HISTORY_NEED, "false")
                                .param("agents_list", agentSb.toString()))
                        .advisors()
                        .call()
                        .entity(new CustomerBeanOutputConverter<>(Dispatcher.DispatcherOut.class));
        Assert.notNull(dispatcherOut, "Dispatch result cannot be null");
        String agentName = dispatcherOut.agent_name;
        String prompt = dispatcherOut.prompt !=null ? dispatcherOut.prompt : messages.get(0).getText();
        if (StringUtils.isBlank(agentName)) {
            throw new IllegalArgumentException("Returned Agent is empty");
        }
        if (!selectionAgents.containsKey(agentName)) {
            throw new IllegalArgumentException("Returned Agent is not in the matching range: " + agentName + ", matching list: " + JsonUtil.toJson(selectionAgents));
        }
        return agentManager.getAgent(agentName).stream(conversationId, List.of(new UserMessage(prompt)),context);
    }

    /**
     * Dispatcher Bean
     *
     * @param agent_name
     */
    record DispatcherOut(String agent_name,String prompt) {

    }


}
